#include <stdio.h>
#include <stdlib.h> 
#include <math.h>
#include <memory.h>
#include <mpi.h>
#define TIMES 150

int size,rank,sizex,sizey,rankx,ranky;
int *idx;
int *idy;
MPI_Status status;
double timeCacl = 0,timeInit = 0,timeRouter = 0;
char *ABIN;
char *BBIN;
char *XBIN;
int X,Y,Z;
double THRESSHOLD;

int getfactor()
{
	int i;
	if(size < 4)
	{
		MPI_Finalize();
		if(rank == 0)
			printf("0 0.000000 0.000000 0.000000 0 0\n");
		exit(0);
	}
	i = sqrt(size);
	if(size % i == 0)
		return i;
	for(i = sqrt(size);i > 1;i--)
	{
		if(size%i == 0)
			return i;
	}
	MPI_Finalize();
	if(rank == 0)
		printf("0.00 0.00 0.00 0 0\n");	
	exit(0);
	return size;
}

void plus(double **a,double *b,double *r,int i,int j,int k)
{
	int YZ = Y*Z;
	int iYZ = i*Y*Z;
	int iYZ1 = (i-1+X)%X*Y*Z;
	int iYZ2 = (i+1)%X*Y*Z;
	int jZ = j*Z;
	int jZ1 = (j-1+Y)%Y*Z;
	int jZ2 = (j+1)%Y*Z;
	int k1 = (k-1+Z)%Z;
	int k2 = (k+1)%Z;
	if(j > 0 && j < Y)
		r[iYZ+jZ+k] = 	a[iYZ+jZ+k][0]*b[iYZ+jZ+k]+
						a[iYZ+jZ+k][1]*b[iYZ1+jZ+k]+
						a[iYZ+jZ+k][2]*b[iYZ2+jZ+k]+
						a[iYZ+jZ+k][3]*b[iYZ+jZ1+k]+
						a[iYZ+jZ+k][4]*b[iYZ+jZ2+k]+
						a[iYZ+jZ+k][5]*b[iYZ2+jZ2+k]+
						a[iYZ+jZ+k][6]*b[iYZ2+jZ1+k]+
						a[iYZ+jZ+k][7]*b[iYZ1+jZ1+k]+
						a[iYZ+jZ+k][8]*b[iYZ1+jZ2+k]+
						a[iYZ+jZ+k][9]*b[iYZ+jZ+k1]+
						a[iYZ+jZ+k][10]*b[iYZ1+jZ+k1]+
						a[iYZ+jZ+k][11]*b[iYZ2+jZ+k1]+
						a[iYZ+jZ+k][12]*b[iYZ+jZ1+k1]+
						a[iYZ+jZ+k][13]*b[iYZ+jZ2+k1]+
						a[iYZ+jZ+k][14]*b[iYZ+jZ+k2]+
						a[iYZ+jZ+k][15]*b[iYZ1+jZ+k2]+
						a[iYZ+jZ+k][16]*b[iYZ2+jZ+k2]+
						a[iYZ+jZ+k][17]*b[iYZ+jZ1+k2]+
						a[iYZ+jZ+k][18]*b[iYZ+jZ2+k2];
	else if(j == 0)
		r[iYZ+jZ+k] = 	a[iYZ+jZ+k][0]*b[iYZ+k]+
						a[iYZ+jZ+k][1]*b[iYZ1+k]+
						a[iYZ+jZ+k][2]*b[iYZ2+k]+
						a[iYZ+jZ+k][4]*b[iYZ+jZ2+k]+
						a[iYZ+jZ+k][5]*b[iYZ2+jZ2+k]+
						a[iYZ+jZ+k][8]*b[iYZ1+jZ2+k]+
						a[iYZ+jZ+k][9]*b[iYZ+jZ+k1]+														
						a[iYZ+jZ+k][10]*b[iYZ1+jZ+k1]+
						a[iYZ+jZ+k][11]*b[iYZ2+jZ+k1]+
						a[iYZ+jZ+k][13]*b[iYZ+jZ2+k1]+
						a[iYZ+jZ+k][14]*b[iYZ+jZ+k2]+
						a[iYZ+jZ+k][15]*b[iYZ1+jZ+k2]+
						a[iYZ+jZ+k][16]*b[iYZ2+jZ+k2]+
						a[iYZ+jZ+k][18]*b[iYZ+jZ2+k2]+						
						a[iYZ+jZ+k][3]*b[(i+X/2)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][6]*b[(i+X/2-1)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][7]*b[(i+X/2+1)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][12]*b[(i+X/2)%X*YZ+jZ+k1]+
						a[iYZ+jZ+k][17]*b[(i+X/2)%X*YZ+jZ+k2];
	else
		r[iYZ+jZ+k] = 	a[iYZ+jZ+k][0]*b[iYZ+jZ+k]+
						a[iYZ+jZ+k][1]*b[iYZ1+jZ+k]+
						a[iYZ+jZ+k][2]*b[iYZ2+jZ+k]+						
						a[iYZ+jZ+k][3]*b[iYZ+jZ1+k]+
						a[iYZ+jZ+k][6]*b[iYZ2+jZ1+k]+
						a[iYZ+jZ+k][7]*b[iYZ1+jZ1+k]+
						a[iYZ+jZ+k][9]*b[iYZ+jZ+k1]+
						a[iYZ+jZ+k][10]*b[iYZ1+jZ+k1]+
						a[iYZ+jZ+k][11]*b[iYZ2+jZ+k1]+
						a[iYZ+jZ+k][12]*b[iYZ+jZ1+k1]+
						a[iYZ+jZ+k][14]*b[iYZ+jZ+k2]+
						a[iYZ+jZ+k][15]*b[iYZ1+jZ+k2]+
						a[iYZ+jZ+k][16]*b[iYZ2+jZ+k2]+
						a[iYZ+jZ+k][17]*b[iYZ+jZ1+k2]+
						a[iYZ+jZ+k][4]*b[(i+X/2)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][5]*b[(i+X/2-1)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][8]*b[(i+X/2+1)%X*YZ+jZ+k]+
						a[iYZ+jZ+k][13]*b[(i+X/2)%X*YZ+jZ+k1]+
						a[iYZ+jZ+k][18]*b[(i+X/2)%X*YZ+jZ+k2];
}

void amultiply(double **a,double *b,double *r)
{			
	int i,j,k;		
	int x1 = idx[rankx];
	int x2 = idx[rankx+1]-1;
	int y1 = idy[ranky];
	int y2 = idy[ranky+1]-1;
	int x0 = (x1-1+X)%X;
	int x3 = (x2+1+X)%X;
	int y0 = (y1-1+Y)%Y;
	int y3 = (y2+1+Y)%Y;
	int rx1 = (rankx-1+sizex)%sizex;
	int rx2 = (rankx+1+sizex)%sizex;
	int ry1 = (ranky-1+sizey)%sizey;
	int ry2 = (ranky+1+sizey)%sizey;
		
	double time1,time2;
	time1 = MPI_Wtime();

	double *st1 = malloc(sizeof(double)*(x2-x1+1)*Z);
	double *rt1 = malloc(sizeof(double)*(x2-x1+1)*Z);

	double *st2 = malloc(sizeof(double)*(y2-y1+3)*Z);
	double *rt2 = malloc(sizeof(double)*(y2-y1+3)*Z);
	
	//纬度
	for(k = 0;k <= x2-x1;k++)
		for(j = 0;j < Z;j++)
			st1[k*Z+j] = b[(k+x1)*Y*Z+y1*Z+j];
	MPI_Sendrecv(st1,(x2-x1+1)*Z,MPI_DOUBLE,rankx*sizey+ry1,rank,
				 rt1,(x2-x1+1)*Z,MPI_DOUBLE,rankx*sizey+ry2,rankx*sizey+ry2,MPI_COMM_WORLD,&status);
	for(k = 0;k <= x2-x1;k++)
		for(j = 0;j < Z;j++)
		{
			b[(k+x1)*Y*Z+y3*Z+j] = rt1[k*Z+j];
			st1[k*Z+j] = b[(k+x1)*Y*Z+y2*Z+j];
		}
	MPI_Sendrecv(st1,(x2-x1+1)*Z,MPI_DOUBLE,rankx*sizey+ry2,rank,
				 rt1,(x2-x1+1)*Z,MPI_DOUBLE,rankx*sizey+ry1,rankx*sizey+ry1,MPI_COMM_WORLD,&status);
	for(k = 0;k <= x2-x1;k++)
		for(j = 0;j < Z;j++)
			b[(k+x1)*Y*Z+y0*Z+j] = rt1[k*Z+j];
	free(st1);
	free(rt1);
	//经度
	for(k = 0;k <= y2-y1+2;k++)
		for(j = 0;j < Z;j++)
			st2[k*Z+j] = b[x1*Y*Z+(k+Y+y1-1)%Y*Z+j];
	MPI_Sendrecv(st2,(y2-y1+3)*Z,MPI_DOUBLE,rx1*sizey+ranky,rank,
				 rt2,(y2-y1+3)*Z,MPI_DOUBLE,rx2*sizey+ranky,rx2*sizey+ranky,MPI_COMM_WORLD,&status);
	for(k = 0;k <= y2-y1+2;k++)
		for(j = 0;j < Z;j++)
			b[x3*Y*Z+(k+Y+y1-1)%Y*Z+j] = rt2[k*Z+j];
	for(k = 0;k <= y2-y1+2;k++)
		for(j = 0;j < Z;j++)			
			st2[k*Z+j] = b[x2*Y*Z+(k+Y+y1-1)%Y*Z+j];
	MPI_Sendrecv(st2,(y2-y1+3)*Z,MPI_DOUBLE,rx2*sizey+ranky,rank,
				 rt2,(y2-y1+3)*Z,MPI_DOUBLE,rx1*sizey+ranky,rx1*sizey+ranky,MPI_COMM_WORLD,&status);
	for(k = 0;k <= y2-y1+2;k++)
		for(j = 0;j < Z;j++)
			b[x0*Y*Z+(k+Y+y1-1)%Y*Z+j] = rt2[k*Z+j];
	free(st2);
	free(rt2);
	
	//北极
	double *t = malloc(sizeof(double)*X*Z);
	if(ranky == 0)
		for(i = x1;i <= x2;i++)
			for(j = 0;j < Z;j++)
				t[i*Z+j] = b[i*Y*Z+j];
	for(i = 1;i < sizex;i++)
	{
		if(rank == i*sizey)
			MPI_Send(&t[idx[i]*Z],(idx[i+1]-idx[i])*Z,MPI_DOUBLE,0,rank,MPI_COMM_WORLD);
		if(rank == 0)		
			MPI_Recv(&t[idx[i]*Z],(idx[i+1]-idx[i])*Z,MPI_DOUBLE,i*sizey,i*sizey,MPI_COMM_WORLD,&status);
	}
	for(i = 1;i < sizex;i++)
	{
		if(rank == i*sizey)
			MPI_Recv(t,X*Z,MPI_DOUBLE,0,rank,MPI_COMM_WORLD,&status);
		if(rank == 0)
			MPI_Send(t,X*Z,MPI_DOUBLE,i*sizey,i*sizey,MPI_COMM_WORLD);
	}
	if(ranky == 0 && rank != 0)
		for(i = 0;i < X;i++)
			for(j = 0;j < Z;j++)
				b[i*Y*Z+j] = t[i*Z+j];	
	//南极
	if(ranky == sizey-1)
		for(i = x1;i <= x2;i++)
			for(j = 0;j < Z;j++)
				t[i*Z+j] = b[i*Y*Z+j];
	for(i = 0;i < sizex-1;i++)
	{
		if(rank == i*sizey+sizey-1)
			MPI_Send(&t[idx[i]*Z],(idx[i+1]-idx[i])*Z,MPI_DOUBLE,size-1,rank,MPI_COMM_WORLD);
		if(rank == size-1)	
			MPI_Recv(&t[idx[i]*Z],(idx[i+1]-idx[i])*Z,MPI_DOUBLE,i*sizey+sizey-1,i*sizey+sizey-1,MPI_COMM_WORLD,&status);
	}
	for(i = 0;i < sizex-1;i++)
	{
		if(rank == i*sizey+sizey-1)
			MPI_Recv(t,X*Z,MPI_DOUBLE,size-1,rank,MPI_COMM_WORLD,&status);
		if(rank == size-1)
			MPI_Send(t,X*Z,MPI_DOUBLE,i*sizey+sizey-1,i*sizey+sizey-1,MPI_COMM_WORLD);
	}
	if(ranky == sizey-1 && rank != size-1)
		for(i = 0;i < X;i++)
			for(j = 0;j < Z;j++)
				b[i*Y*Z+j] = t[i*Z+j];
	free(t);
	time2 = MPI_Wtime();
	timeRouter += time2-time1;
	
	for(i = idx[rankx];i < idx[rankx+1];i++)
		for(j = idy[ranky];j < idy[ranky+1];j++)
			for(k = 0;k < Z;k++)
				plus(a,b,r,i,j,k);
}

int stop(double **a,double *x,double *b)
{	
	double *t = malloc(sizeof(double)*X*Y*Z);
	amultiply(a,x,t);
	int i,j,k;

	double all = 0;
	for(i = idx[rankx];i < idx[rankx+1];i++)
		for(j = idy[ranky];j < idy[ranky+1];j++)
			for(k = 0;k < Z;k++)
				all += pow(b[i*Y*Z+j*Z+k]-t[i*Y*Z+j*Z+k],2);
	free(t);
	double tall;	
	double time1,time2;
	time1 = MPI_Wtime();	
	MPI_Allreduce(&all,&tall,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);	
	time2 = MPI_Wtime();
	timeRouter += time2-time1;
	tall = sqrt(tall);
	
	//if(rank == 0)
	//	printf("%.15lf\n",tall);
	if(tall < THRESSHOLD)
		return 1;
	return 0;
}

double dotproduct(double *x,double *y)
{
	int i,j,k;
	double all = 0;
	for(i = idx[rankx];i < idx[rankx+1];i++)
		for(j = idy[ranky];j < idy[ranky+1];j++)
			for(k = 0;k < Z;k++)
				all += x[i*Y*Z+j*Z+k]*y[i*Y*Z+j*Z+k];
	double tall;
	double time1,time2;
	time1 = MPI_Wtime();	
	MPI_Allreduce(&all,&tall,1,MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD);
	time2 = MPI_Wtime();
	timeRouter += time2-time1;
	return tall;
}

int main(int argc, char *argv[])
{	
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

	double time1,time2;
	FILE *fr;
	int i,j,k;
	
	time1 = MPI_Wtime();
	sizex = 0;
	sizey = 0;
	if(argc == 1)
	{
		ABIN = "case1/A.bin";
		BBIN = "case1/b.bin";
		XBIN = "case1/x0.bin";
		X = 360;
		Y = 180;
		Z = 38;
		THRESSHOLD = 1.878932301845534E-005;		
	}
	else if(argc == 2)
	{
		int t;
		sscanf(argv[1],"%d",&t);
		if(t == 1)
		{
			ABIN = "case1/A.bin";
			BBIN = "case1/b.bin";
			XBIN = "case1/x0.bin";
			X = 360;
			Y = 180;
			Z = 38;
			THRESSHOLD = 1.878932301845534E-005;
		}
		else if(t == 2)
		{
			ABIN = "case2/A.bin";
			BBIN = "case2/b.bin";
			XBIN = "case2/x0.bin";
			X = 720;
			Y = 360;
			Z = 38;
			THRESSHOLD = 6.724779532126749E-005;
		}
		else
			exit(0);
	}	
	else if(argc == 4)
	{
		sscanf(argv[2],"%d",&sizex);
		sscanf(argv[3],"%d",&sizey);		
	}
	else
		exit(0);
	if(sizex <= 1 || sizey <= 1)
	{		
		sizey = getfactor();
		sizex = size/sizey;
	}
		
	idx = malloc(sizeof(int)*(sizex+1));
	for(i = 0;i < sizex;i++)
		idx[i] = X/sizex*i;
	idx[sizex] = X;
	
	idy = malloc(sizeof(int)*(sizey+1));
	for(i = 0;i < sizey;i++)
		idy[i] = Y/sizey*i;
	idy[sizey] = Y;

	rankx = rank/sizey;
	ranky = rank%sizey;

	double **a = malloc(sizeof(double*)*X*Y*Z);
	double *b = malloc(sizeof(double)*X*Y*Z);
	double *x = malloc(sizeof(double)*X*Y*Z);
	for(i = 0;i < X*Y*Z;i++)
		a[i] = malloc(sizeof(double)*19);
	
	//IO
	int fstart,fend;	
	fstart = idx[rankx]*Y*Z;
	fend = idx[rankx+1]*Y*Z;
	
	fr = fopen(ABIN,"rb");
	fseek(fr,fstart*8*19,SEEK_SET);
	for(i = fstart;i < fend;i++)
		if(fread(a[i],sizeof(double),19,fr) != 19)
			return 1;
	fclose(fr);	
	fr = fopen(BBIN,"rb");
	fseek(fr,fstart*8,SEEK_SET);
	for(i = fstart;i < fend;i++)
		if(fread(&b[i],sizeof(double),1,fr) != 1)
			return 1;
	fclose(fr);	
	fr = fopen(XBIN,"rb");
	fseek(fr,fstart*8,SEEK_SET);
	for(i = fstart;i < fend;i++)
		if(fread(&x[i],sizeof(double),1,fr) != 1)
			return 1;
	fclose(fr);
	
	//Init
	double *r = malloc(sizeof(double)*X*Y*Z);
	double *t = malloc(sizeof(double)*X*Y*Z);
	double *p = malloc(sizeof(double)*X*Y*Z);
	double *q = malloc(sizeof(double)*X*Y*Z);
	double alpha,beta,gama;

	time2 = MPI_Wtime();
	timeInit = time2-time1;

	//Calc

	time1 = MPI_Wtime();

	amultiply(a,x,r);
	
	for(i = idx[rankx];i < idx[rankx+1];i++)
		for(j = idy[ranky];j < idy[ranky+1];j++)
			for(k = 0;k < Z;k++)
				r[i*Y*Z+j*Z+k] = b[i*Y*Z+j*Z+k]-r[i*Y*Z+j*Z+k];	
	
	memcpy(p,r,sizeof(double)*X*Y*Z);	
	amultiply(a,p,q);

	int round = 0;
	
	//if(rank == 0)
		//printf("%d %d\n",sizex,sizey);
		
	while(round < TIMES)
	{				
		gama = dotproduct(q,q);
		alpha = dotproduct(r,q)/gama;

		for(i = idx[rankx];i < idx[rankx+1];i++)
			for(j = idy[ranky];j < idy[ranky+1];j++)
				for(k = 0;k < Z;k++)
					x[i*Y*Z+j*Z+k] += alpha*p[i*Y*Z+j*Z+k];		
		if(stop(a,x,b))
			break;
		for(i = idx[rankx];i < idx[rankx+1];i++)
			for(j = idy[ranky];j < idy[ranky+1];j++)
				for(k = 0;k < Z;k++)
					r[i*Y*Z+j*Z+k] -= alpha*q[i*Y*Z+j*Z+k];				
		amultiply(a,r,t);
		beta = -dotproduct(t,q)/gama;	
		for(i = idx[rankx];i < idx[rankx+1];i++)
			for(j = idy[ranky];j < idy[ranky+1];j++)
				for(k = 0;k < Z;k++)
					p[i*Y*Z+j*Z+k] = r[i*Y*Z+j*Z+k]+beta*p[i*Y*Z+j*Z+k];			
		for(i = idx[rankx];i < idx[rankx+1];i++)
			for(j = idy[ranky];j < idy[ranky+1];j++)
				for(k = 0;k < Z;k++)
					q[i*Y*Z+j*Z+k] = t[i*Y*Z+j*Z+k]+beta*q[i*Y*Z+j*Z+k];
		round++;
	}
	time2 = MPI_Wtime();
	timeCacl = time2-time1-timeRouter;
	if(rank == 0)
		printf("%d %lf %lf %lf %d %d\n",round,timeInit,timeCacl,timeRouter,sizex,sizey);
	MPI_Finalize();
	return 0;
}
